package multi_threader

import (
	"context"
	"fmt"
	"runtime"
	"strconv"
	"sync"
)

// Pool 线程池
//
//	threaderNum 开启线程数量,范围1~可用的逻辑 CPU 数
//	inChan 输入通道
//	outChan 输出通道
//	arg 用户配置
type Pool struct {
	threaderNum int
	inChan      chan any
	outChan     chan any
	producerWg  sync.WaitGroup
	consumerWg  sync.WaitGroup
	ctx         context.Context
	cancel      context.CancelFunc
	isPrint     bool
}

// NewPool 创建线程池
//
//	threaderNum 开启线程数量,范围1~可用的逻辑 CPU 数
//	arg 用户配置
func NewPool(ctx context.Context, threaderNum int, isPrint bool) *Pool {
	pool := &Pool{
		threaderNum: threaderNum,
		inChan:      make(chan any, threaderNum*2),
		outChan:     make(chan any, threaderNum*2),
		producerWg:  sync.WaitGroup{},
		consumerWg:  sync.WaitGroup{},
		isPrint:     isPrint,
	}
	// 判断设置开启线程数
	if pool.threaderNum <= 0 {
		pool.threaderNum = 1
	} else if pool.threaderNum > runtime.NumCPU() {
		pool.threaderNum = runtime.NumCPU()
	}
	// 上下文控制
	if ctx == nil {
		ctx = context.Background()
	}
	pool.ctx, pool.cancel = context.WithCancel(ctx)

	return pool
}

// EnablePrint 设置是否输出日志内容
func (p *Pool) EnablePrint(enable bool) *Pool {
	p.isPrint = enable
	return p
}

// Run 运行
//
//	productFun 生产者函数 返回待处理数据集切片
//	consumerFun 消费者函数，入参为索引和处理者返回数据
//	handleFun 处理者函数，入参为生产者返回数据集中的每个元素
func (p *Pool) Run(_func poolFunc) {

	// 设置go调用cpu核心数，新版本默认最大核心
	runtime.GOMAXPROCS(runtime.NumCPU())

	// 开启生产者线程
	p.producerWg.Add(1)
	go p.producer(_func.ProductFunc)

	// 开启执行者线程
	for i := 0; i < p.threaderNum; i++ {
		p.producerWg.Add(1)
		go p.handle(i, _func.HandleFunc)
	}

	// 开启消费者线程，记录结果
	p.consumerWg.Add(1)
	go p.consumer(_func.ConsumerFunc)

	// 等待生产者、执行者线程结束
	p.producerWg.Wait()

	// 执行者结束后,关闭消费者通道
	close(p.outChan)

	// 等待消费者线程结束
	p.consumerWg.Wait()

}

// Stop 停止方法
func (p *Pool) Stop() {
	p.cancel() // 通知所有协程停止
	close(p.inChan)
	p.producerWg.Wait()
	close(p.outChan)
	p.consumerWg.Wait()
}

// 生产者
/*
producer 启动一个生产者协程，负责调用指定的生产函数将数据发送到通道。

参数:
  f   生产函数，接收 *chan any 类型参数并返回 []any 数据切片。该函数由调用者提供，用于生成需要发送的数据。

说明:
  - 持续监听上下文取消信号，一旦收到信号立即关闭通道并退出循环。
  - 在默认情况下，调用生产函数获取数据，并逐个发送到 inChan 通道中。
  - 数据发送完成后关闭 inChan 通道并退出当前协程。
*/
func (p *Pool) producer(f func() ([]any, <-chan any)) {
	defer p.producerWg.Done()
	if p.isPrint {
		fmt.Println("info", "生产者线程启动...")
		defer fmt.Println("info", "生产者线程结束")
	}

	for {
		select {
		case <-p.ctx.Done():
			// 收到上下文取消信号后关闭通道，停止生产流程
			close(p.inChan)
			return
		default:
			// 调用生产函数获取待发送的数据
			sourceData, _inChan := f()
			if sourceData != nil {
				// 遍历数据集，依次向通道发送每个元素
				for _, item := range sourceData {
					p.inChan <- item
				}
			} else if _inChan != nil {
				// 从通道中读取数据，依次向输入通道发送每个元素
				for item := range _inChan {
					p.inChan <- item
				}
			}

			// 数据全部发送完成后关闭通道，通知消费者无更多数据
			close(p.inChan)
			return
		}
	}
}

// 处理者
/*
handle 处理协程，负责从通道中消费数据并调用指定处理函数进行处理。

参数:
  threadId 当前线程的唯一标识，用于日志打印和调试。
  f       处理函数，接收 any 类型参数并返回 any 类型结果。该函数由调用者提供，用于处理输入数据。

说明:
  - 持续监听上下文取消信号，一旦收到信号立即退出协程。
  - 从 inChan 通道中持续消费数据，对每个元素调用处理函数 f。
  - 处理结果通过 outChan 通道发送出去。
*/
func (p *Pool) handle(threadId int, f func(any) any) {
	defer p.producerWg.Done()
	if p.isPrint {
		fmt.Println("info", "开启线程id="+strconv.Itoa(threadId))
		defer fmt.Println("info", "线程id="+strconv.Itoa(threadId)+"结束")
	}

	// 处理线程主循环，持续从 inChan 中读取数据
	for item := range p.inChan {
		select {
		case <-p.ctx.Done():
			// 收到上下文取消信号，立即退出协程
			return
		default:
			// 对当前数据项执行处理函数，并将结果发送至 outChan
			p.outChan <- f(item)
		}
	}
}

// 消费者
/*
consumer 启动一个消费者协程，负责从输出通道中消费数据并调用指定的消费函数进行处理。

参数:
  f   消费函数，接收 int 和 any 类型参数，无返回值。该函数由调用者提供，用于处理从 outChan 接收到的数据项。
      第一个参数为数据项的序号（从 0 开始递增），第二个参数为实际的数据项。

说明:
  - 从 outChan 通道中持续读取数据，直到通道被关闭。
  - 每次读取到数据后调用消费函数 f，并传入当前项的序号和数据本身。
*/
func (p *Pool) consumer(f func(int, any)) {
	defer p.consumerWg.Done()
	if p.isPrint {
		fmt.Println("info", "消费者线程启动")
		defer fmt.Println("info", "消费者线程结束")
	}

	// 从输出通道中持续取出数据并处理
	n := 0
	for item := range p.outChan {
		// 调用消费函数处理当前数据项，并递增序号
		f(n, item)
		n += 1
	}
}

// poolFunc 处理函数接口
type poolFunc interface {
	// ProductFunc 生产者执行程序
	// 输入参数为通道指针, 返回值为待处理数据列表
	// 两者任选其一
	ProductFunc() ([]any, <-chan any)
	// ConsumerFunc 消费者处理程序，入参为处理序号号处理线程输出结果
	ConsumerFunc(int, any)
	// HandleFunc 处理者处理程序，入参为生产者列表中的元素，输出为处理后结果
	HandleFunc(any) any
}
